from __future__ import division
import os
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import xml.dom.minidom
import numpy as np
import threading
from concurrent.futures import ThreadPoolExecutor

# 指定线程的数量
num_threads = 32

ImgPath = 'SN002_20231103_960_540'
AnnoPath = 'SN002_20231103_960_540_xml_one_class'
ProcessedPath = 'SN002_20231103_960_540_xml_one_class_crop'
if not os.path.exists(ProcessedPath):
    os.mkdir(ProcessedPath)

filelist = os.listdir(ImgPath)
lock = threading.Lock()

def process_image(file):
    if file.endswith('.jpg'):
        image_pre, ext = os.path.splitext(file)
        imgfile = ImgPath + '/' + file
        print(imgfile)
        if not os.path.exists(AnnoPath + '/' + image_pre + '.xml' ):
            return
        xmlfile = AnnoPath + '/' + image_pre + '.xml'
        DomTree = xml.dom.minidom.parse(xmlfile)
        annotation = DomTree.documentElement
        filenamelist = annotation.getElementsByTagName('filename')
        filename = filenamelist[0].childNodes[0].data
        objectlist = annotation.getElementsByTagName('object')
        i = 1
        for objects in objectlist:
            namelist = objects.getElementsByTagName('name')
            objectname = namelist[0].childNodes[0].data
            savepath = objectname        
            if not os.path.exists(ProcessedPath + '/' + savepath):
                os.makedirs(ProcessedPath + '/' + savepath)
            bndbox = objects.getElementsByTagName('bndbox')
            cropboxes = []
            for box in bndbox:
                x1_list = box.getElementsByTagName('xmin')
                x1 = float(x1_list[0].childNodes[0].data)
                y1_list = box.getElementsByTagName('ymin')
                y1 = float(y1_list[0].childNodes[0].data)
                x2_list = box.getElementsByTagName('xmax')
                x2 = float(x2_list[0].childNodes[0].data)
                y2_list = box.getElementsByTagName('ymax')
                y2 = float(y2_list[0].childNodes[0].data)
                w = x2 - x1
                h = y2 - y1
                obj = np.array([x1,y1,x2,y2])
                shift = np.array([[1,1,1,1]])
                XYmatrix = np.tile(obj,(1,1))
                cropboxes = XYmatrix * shift
                img = Image.open(imgfile)
                for cropbox in cropboxes:
                    with lock:
                        print("Cropbox:", cropbox)
                        cropedimg = img.crop(cropbox)
                        if not cropedimg:
                            print("Warning: Empty cropedimg")
                        else:
                            cropedimg.save((ProcessedPath + '/' + savepath + '/' + image_pre + '_' + str(i) + '.jpg'), quality=95, subsampling=0)
                        i += 1


# 创建一个线程池
with ThreadPoolExecutor(max_workers=num_threads) as executor:
    # 使用线程池来处理图像
    executor.map(process_image, filelist)

print("Processing completed.")
